library(tidyverse)
library(maptools)
library(parallel)
library(redist)
library(spdep)

ipw <- function(x, beta, pop){
    indpop <- which(x$distance_parity <= pop)
    indbeta <- which(x$beta_sequence == beta)
    inds <- intersect(indpop, indbeta)
    psi <- x$constraint_pop[inds]
    w <- 1 / exp(beta * psi)
    inds <- sample(inds, length(inds), replace = TRUE, prob = w)
    x <- x$partitions[,inds]
    return(x)
}
ben_theme <- function(){
    theme_classic() +
        theme(panel.background = element_blank(),
              panel.grid.major = element_blank(),
              panel.grid.minor = element_blank(),
              axis.line = element_line(colour = "black"),
              panel.border = element_rect(colour = "black", fill = NA, size = 1),
              strip.background = element_blank(),
              legend.position = "bottom", legend.title = element_blank(),
              plot.title = element_text(hjust = 0.5),
              plot.subtitle = element_text(hjust = .5))
}

## Get info on job
i <- as.numeric(Sys.getenv("SLURM_ARRAY_TASK_ID")) 
set.seed(i)
## -------------------
## Load data and clean
## -------------------
ia <- readShapePoly("../../data/ia/county.shp")
votes <- read_csv("../../data/ia/ia_vote.csv")
pops <- read_csv("../../data/ia/ia_pop.csv") %>%
    mutate(County = gsub(" County", "", County))
info <- inner_join(votes, pops) %>%
    rename(county = County,
           trump = `Trump #`,
           clinton = `Clinton #`,
           other = `Others #`,
           total = `Total`,
           fips = `FIPS code[10]`,
           pop = Population) %>%
    select(county, trump, clinton, fips, pop)

## Join with shapefile
ia@data$COUNTY <- as.character(ia@data$COUNTY)
ia@data$COUNTY[ia@data$COUNTY == "Obrien"] <- "O'Brien"
ia@data$id <- 1:nrow(ia@data)
ia@data <- inner_join(ia@data, info, by = c("COUNTY" = "county"))

## --------------
## Run redist.rsg
## --------------
adjlist <- poly2nb(ia, queen = FALSE)
adjlist <- lapply(adjlist, function(x){x-1})
nsims <- 10000

## Unconstrained
sv_out <- lapply(1:nsims, function(x){
    if(x %% (nsims / 10) == 0){
        print(paste0("Done with ", x, " no parity iterations at ", Sys.time(), "."))
    }
    return(redist.rsg(
        adjlist, ia@data$pop, ndists = 4, thresh = 100, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
})
rsg_out <- do.call(cbind, sv_out)
rsg_seg_full <- redist.segcalc(rsg_out, grouppop = ia@data$trump, 
                               fullpop = ia@data$pop)
rsg_out_full <- rsg_out

## 5%
sv_out <- lapply(1:nsims, function(x){
    if(x %% (nsims / 10) == 0){
        print(paste0("Done with ", x, " 5% iterations at ", Sys.time(), "."))
    }
    return(redist.rsg(
        adjlist, ia@data$pop, ndists = 4, thresh = .05, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
})
rsg_out <- do.call(cbind, sv_out)
rsg_seg_05p <- redist.segcalc(rsg_out, grouppop = ia@data$trump, 
                              fullpop = ia@data$pop)

rsg_out_5 <- rsg_out

## 1%
sv_out <- lapply(1:nsims, function(x){
    if(x %% (nsims / 10) == 0){
        print(paste0("Done with ", x, " 1% iterations at ", Sys.time(), "."))
    }
    return(redist.rsg(
        adjlist, ia@data$pop, ndists = 4, thresh = .01, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
})
rsg_out <- do.call(cbind, sv_out)
rsg_seg_01p <- redist.segcalc(rsg_out, grouppop = ia@data$trump, 
                              fullpop = ia@data$pop)

rsg_out_1 <- rsg_out

## -------------
## Create output
## -------------
df_out <- data.frame(seg = c(rsg_seg_full, rsg_seg_05p, rsg_seg_01p),
                     parity = c(rep("No Equal Population Constraint", nsims),
                                rep("5% Equal Population Constraint", nsims),
                                rep("1% Equal Population Constraint", nsims)))
write_csv(df_out, path = paste0("/n/imai_lab/ckenny/rsg_", i, ".csv"))

df_districts <- do.call(cbind, list(rsg_out_full,rsg_out_5,rsg_out_1))
saveRDS(object = df_districts, file = paste0("/n/imai_lab/ckenny/rsg/rsg_districts_", i, ".csv"))



